import pymysql


class MysqlHelper():
    def __init__(self, host, database, user, password, port=3306, charset="utf8"):
        self.host = host
        self.database = database
        self.user = user
        self.password = password
        self.port = port
        self.charset = charset

    def connect(self):
        self.conn = pymysql.connect(
            host=self.host,
            database=self.database,
            user=self.user,
            password=self.password,
            port=self.port,
            charset=self.charset
        )
        self.cursor = self.conn.cursor()

    def close(self):
        self.cursor.close()
        self.conn.close()

    def insert(self, sql, params=[]):
        try:
            self.connect()
            insert_count = self.cursor.execute(sql, params)  # 获取影响的行数
            self.conn.commit()
            self.close()
            return insert_count
        except Exception as e:
            print(f"数据表插入错误: {e}")
            return 0

    def select_one(self, sql, params=[]):
        try:
            self.connect()
            self.cursor.execute(sql, params)
            find_count = self.cursor.fetchone()
            self.close()
            return find_count
        except Exception as e:
            print(f"数据表查询错误: {e}")
            return (0,)  # 返回包含0的元组


class Person():
    def register(self):
        name = input("请输入注册用户名：")
        pwd = input("请输入注册密码：")
        helper = MysqlHelper("127.0.0.1", "hqyj", "root", "root")
        sql = "insert into tb_users(username, password) values (%s, %s)"
        ret = helper.insert(sql, [name, pwd])
        if ret > 0:  # 直接判断返回值
            print("注册成功")
        else:
            print("注册失败")

    def login(self):
        name = input("请输入登录用户名：")
        pwd = input("请输入登录密码：")
        helper = MysqlHelper("127.0.0.1", "hqyj", "root", "root")
        sql = "select count(*) from tb_users where username=%s and password=%s;"  # 修正字段名
        ret = helper.select_one(sql, [name, pwd])
        if ret[0] > 0:  # ret是元组，取第一个元素
            print("登录成功")
        else:
            print("登录失败，请重新登录")


def main():
    print("---淘宝页面---\n"
          "1、注册\n"
          "2、登录\n"
          "3、退出")
    user = Person()
    while True:
        choice = input("请输入要进行的操作:")
        if choice == "1":
            user.register()
        elif choice == "2":
            user.login()
        elif choice == "3":
            break
        else:
            print("输入有误，请重新输入！！！")


if __name__ == "__main__":
    main()
